#include <iostream>
#include <string>
#include <vector>
#include <chrono>
#include <math.h>

/**
 * intel SIMD Docs: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
 */


#pragma GCC optimize(0)
using namespace std;
using namespace std::chrono;
#ifdef __SSE2__

#include <emmintrin.h>
#include <immintrin.h>
#endif

struct SqrtImpl {
public:
    /**
     *  Use SSE instruction to compute sqrt(a^2 + b^2) + 0.5
     * @param pArray1
     * @param pArray2
     * @param pResult
     * @param nSize
     */
    static void computeArrayCPlusPlusSSE(
            float *pArray1,                   // [in] first source array
            float *pArray2,                   // [in] second source array
            float *pResult,                   // [out] result array
            int nSize)                        // [in] size of all arrays
    {
        int nLoop = nSize / 4;

        __m128 m1, m2, m3, m4;

        __m128 *pSrc1 = (__m128 *) pArray1;
        __m128 *pSrc2 = (__m128 *) pArray2;
        __m128 *pDest = (__m128 *) pResult;


        __m128 m0_5 = _mm_set_ps1(0.5f);        // m0_5[0, 1, 2, 3] = 0.5

        for (int i = 0; i < nLoop; i++) {
            m1 = _mm_mul_ps(*pSrc1, *pSrc1);        // m1 = *pSrc1 * *pSrc1
            m2 = _mm_mul_ps(*pSrc2, *pSrc2);        // m2 = *pSrc2 * *pSrc2
            m3 = _mm_add_ps(m1, m2);                // m3 = m1 + m2
            m4 = _mm_sqrt_ps(m3);                   // m4 = sqrt(m3)
            *pDest = _mm_add_ps(m4, m0_5);          // *pDest = m4 + 0.5

            pSrc1++;
            pSrc2++;
            pDest++;
        }
    }
    static void computeArrayCPlusPlusAVX(
            float *pArray1,                   // [in] first source array
            float *pArray2,                   // [in] second source array
            float *pResult,                   // [out] result array
            int nSize)                        // [in] size of all arrays
    {
        int nLoop = nSize / 8;

        __m256 m1, m2, m3, m4;

        __m256 *pSrc1 = (__m256 *) pArray1;
        __m256 *pSrc2 = (__m256 *) pArray2;
        __m256 *pDest = (__m256 *) pResult;
//        cout << "avx load finished " << endl;

        __m256 m0_5 = _mm256_set_ps(0.5f,0.5f,0.5f,0.5f,0.5f,0.5f,0.5f,0.5f);        // m0_5[0, 1, 2, 3] = 0.5

        for (int i = 0; i < nLoop; i++) {
            m1 = _mm256_mul_ps(*pSrc1, *pSrc1);        // m1 = *pSrc1 * *pSrc1
            m2 = _mm256_mul_ps(*pSrc2, *pSrc2);        // m2 = *pSrc2 * *pSrc2
            m3 = _mm256_add_ps(m1, m2);                // m3 = m1 + m2
            m4 = _mm256_sqrt_ps(m3);                   // m4 = sqrt(m3)
            *pDest = _mm256_add_ps(m4, m0_5);          // *pDest = m4 + 0.5

            pSrc1++;
            pSrc2++;
            pDest++;
        }
    }

    static void computeArrayCPlusPlus(
            float *pArray1,                   // [in] first source array
            float *pArray2,                   // [in] second source array
            float *pResult,                   // [out] result array
            int nSize)                        // [in] size of all arrays
    {

        int i;

        float *pSource1 = pArray1;
        float *pSource2 = pArray2;
        float *pDest = pResult;

        for (i = 0; i < nSize; i++) {
            *pDest = (float) sqrt((*pSource1) * (*pSource1) + (*pSource2)
                                                              * (*pSource2)) + 0.5f;

            pSource1++;
            pSource2++;
            pDest++;
        }
    }


};


int main() {

    int arraySzie = 10000;
    float pArray1[arraySzie];                 // [in] first source array
    float pArray2[arraySzie];                   // [in] second source array
    float pResult[arraySzie];

    for (int i = 0; i < 1000; i++) {
        pArray1[i] = 1000.0f + i;
        pArray2[i] = 1000.0f + i;
    }

    auto start0 = system_clock::now();
    SqrtImpl::computeArrayCPlusPlusAVX(pArray1, pArray2, pResult, arraySzie);
    auto end0 = system_clock::now();
    auto duration0 = duration_cast<nanoseconds>(end0 - start0);
    cout << "AVX Time cost: " << duration0.count() / 1000.0 << " ms" << endl;


    auto start1 = system_clock::now();
    SqrtImpl::computeArrayCPlusPlusSSE(pArray1, pArray2, pResult, arraySzie);
    auto end1 = system_clock::now();
    auto duration1 = duration_cast<nanoseconds>(end1 - start1);
    cout << "SSE Time cost: " << duration1.count() / 1000.0 << " ms" << endl;


    auto start2 = system_clock::now();
    SqrtImpl::computeArrayCPlusPlus(pArray1, pArray2, pResult, arraySzie);
    auto end2 = system_clock::now();
    auto duration2 = duration_cast<nanoseconds>(end2 - start2);
    cout << "Time cost: " << duration2.count() / 1000.0 << " ms" << endl;

    /*execute result may be :
     * AVX Time cost: 12.165 ms
     * SSE Time cost: 19.564 ms
     * Time cost: 55.863 ms
     *
     */

}